"""tensor切分"""

import torch


def test_chunk():
    print("==== chunk ====")
    a = torch.rand(3, 4)
    print("a: ")
    print(a)

    print("torch.chunk(a, 2, dim=0): ")
    out = torch.chunk(a, 2, dim=0)
    print(out[0], out[0].shape)
    print(out[1], out[1].shape)
    print()


def test_split():
    print("==== split ====")
    a = torch.rand(3, 4)
    print("a: ")
    print(a)

    print("torch.split(a, 2, dim=0): (a=（3，4）--> 3/2 =1余1，故是2个tensor)")
    out = torch.split(a, 2, dim=0)
    print(out[0], out[0].shape)
    print(out[1], out[1].shape)
    print()

    print("torch.split(a, 3, dim=1): (a=（3，4）--> 4/3 =1余1，故是2个tensor)")
    out = torch.split(a, 3, dim=1)
    print(out[0], out[0].shape)
    print(out[1], out[1].shape)
    print()

    print("torch.split(a, [1,2,1], dim=1): (将dim=1上的4拆成1，2，1, 故是3个tensor)")
    out = torch.split(a, [1, 2, 1], dim=1)
    print(out[0], out[0].shape)
    print(out[1], out[1].shape)
    print(out[2], out[2].shape)
    print()


if __name__ == '__main__':
    test_chunk()
    test_split()
